import os
os.environ["CUDA_VISIBLE_DEVICES"] ="0"

#Model Download
from modelscope import snapshot_download
#emb_model_dir = snapshot_download('maidalun/bce-embedding-base_v1',cache_dir="D:\\evns\\models\\bce")
#rerank_model_dir =snapshot_download('maidalun/bce-reranker-base_v1',cache_dir="D:\\evns\\models\\bce")

from BCEmbedding import EmbeddingModel
sentences = ['今天天气不错哟','明天一起去徒步']
model = EmbeddingModel(model_name_or_path="D:/evns/models/bce/maidalun/bce-embedding-base_v1")
embeddings = model.encode(sentences)
print(embeddings.shape)


from BCEmbedding import  RerankerModel

query = "一个女人站在高崖上单腿站立，俯瞰一条河流。"
passages =[ "一个女人站在悬崖上。", "一个孩子在她的卧室里读书。"]

sentence_pairs = [[query, passage] for passage in passages]

rerank_model = RerankerModel(model_name_or_path="D:/evns/models/bce/maidalun/bce-reranker-base_v1")

scores = rerank_model.compute_score(sentence_pairs)

rerank_results = rerank_model.rerank(query, passages)

print(rerank_results)
